use async_trait::async_trait;
use rustcloud_core::{ServiceResult, ServiceError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH, Duration};
use tokio::sync::RwLock;
use uuid::Uuid;

/// Span类型枚举
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SpanKind {
    /// 客户端发起的调用
    Client,
    /// 服务端接收的请求
    Server,
    /// 消息生产者
    Producer,
    /// 消息消费者
    Consumer,
    /// 内部操作
    Internal,
}

/// Span状态枚举
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SpanStatus {
    /// 未设置状态
    Unset,
    /// 成功
    Ok,
    /// 错误
    Error,
}

/// 标签类型定义
pub type Tags = HashMap<String, String>;

/// Span属性
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpanAttributes {
    /// 标签
    pub tags: Tags,
    /// 日志
    pub logs: Vec<SpanLog>,
    /// 状态
    pub status: SpanStatus,
    /// 状态消息
    pub status_message: Option<String>,
}

impl Default for SpanAttributes {
    fn default() -> Self {
        Self {
            tags: HashMap::new(),
            logs: Vec::new(),
            status: SpanStatus::Unset,
            status_message: None,
        }
    }
}

/// Span日志条目
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpanLog {
    /// 时间戳（微秒）
    pub timestamp: u64,
    /// 日志级别
    pub level: String,
    /// 日志消息
    pub message: String,
    /// 附加字段
    pub fields: HashMap<String, String>,
}

impl SpanLog {
    pub fn new(level: &str, message: &str) -> Self {
        Self {
            timestamp: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as u64,
            level: level.to_string(),
            message: message.to_string(),
            fields: HashMap::new(),
        }
    }

    pub fn with_field(mut self, key: &str, value: &str) -> Self {
        self.fields.insert(key.to_string(), value.to_string());
        self
    }
}

/// Span引用类型
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SpanReferenceType {
    /// 子Span（最常见的父子关系）
    ChildOf,
    /// 跟随关系（异步操作的关系）
    FollowsFrom,
}

/// Span引用
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpanReference {
    /// 引用类型
    pub reference_type: SpanReferenceType,
    /// 引用的Trace ID
    pub trace_id: String,
    /// 引用的Span ID
    pub span_id: String,
}

/// 分布式追踪的Span
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Span {
    /// Trace ID（整个请求链路的唯一标识）
    pub trace_id: String,
    /// Span ID（当前操作的唯一标识）
    pub span_id: String,
    /// 父Span ID
    pub parent_id: Option<String>,
    /// 操作名称
    pub operation_name: String,
    /// Span类型
    pub kind: SpanKind,
    /// 开始时间（微秒时间戳）
    pub start_time: u64,
    /// 结束时间（微秒时间戳）
    pub end_time: Option<u64>,
    /// 持续时间（微秒）
    pub duration: Option<u64>,
    /// 服务名称
    pub service_name: String,
    /// 资源名称（如HTTP路径、数据库表名等）
    pub resource: String,
    /// Span属性
    pub attributes: SpanAttributes,
    /// Span引用
    pub references: Vec<SpanReference>,
}

impl Span {
    /// 创建新的根Span
    pub fn new_root(operation_name: &str, service_name: &str) -> Self {
        let trace_id = Uuid::new_v4().to_string();
        let span_id = Uuid::new_v4().to_string();
        
        Self {
            trace_id,
            span_id,
            parent_id: None,
            operation_name: operation_name.to_string(),
            kind: SpanKind::Internal,
            start_time: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as u64,
            end_time: None,
            duration: None,
            service_name: service_name.to_string(),
            resource: operation_name.to_string(),
            attributes: SpanAttributes::default(),
            references: Vec::new(),
        }
    }

    /// 创建子Span
    pub fn new_child(&self, operation_name: &str, service_name: &str) -> Self {
        let span_id = Uuid::new_v4().to_string();
        
        Self {
            trace_id: self.trace_id.clone(),
            span_id,
            parent_id: Some(self.span_id.clone()),
            operation_name: operation_name.to_string(),
            kind: SpanKind::Internal,
            start_time: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as u64,
            end_time: None,
            duration: None,
            service_name: service_name.to_string(),
            resource: operation_name.to_string(),
            attributes: SpanAttributes::default(),
            references: vec![SpanReference {
                reference_type: SpanReferenceType::ChildOf,
                trace_id: self.trace_id.clone(),
                span_id: self.span_id.clone(),
            }],
        }
    }

    /// 设置Span类型
    pub fn with_kind(mut self, kind: SpanKind) -> Self {
        self.kind = kind;
        self
    }

    /// 设置资源名称
    pub fn with_resource(mut self, resource: &str) -> Self {
        self.resource = resource.to_string();
        self
    }

    /// 添加标签
    pub fn set_tag(&mut self, key: &str, value: &str) {
        self.attributes.tags.insert(key.to_string(), value.to_string());
    }

    /// 添加日志
    pub fn log(&mut self, level: &str, message: &str) {
        self.attributes.logs.push(SpanLog::new(level, message));
    }

    /// 添加带字段的日志
    pub fn log_with_fields(&mut self, level: &str, message: &str, fields: HashMap<String, String>) {
        let mut log = SpanLog::new(level, message);
        log.fields = fields;
        self.attributes.logs.push(log);
    }

    /// 设置状态
    pub fn set_status(&mut self, status: SpanStatus, message: Option<String>) {
        self.attributes.status = status;
        self.attributes.status_message = message;
    }

    /// 设置错误状态
    pub fn set_error(&mut self, error: &str) {
        self.set_status(SpanStatus::Error, Some(error.to_string()));
        self.set_tag("error", "true");
        self.set_tag("error.message", error);
    }

    /// 完成Span
    pub fn finish(mut self) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_micros() as u64;
        
        self.end_time = Some(now);
        self.duration = Some(now - self.start_time);
        
        if self.attributes.status == SpanStatus::Unset {
            self.attributes.status = SpanStatus::Ok;
        }
        
        self
    }

    /// 判断Span是否已完成
    pub fn is_finished(&self) -> bool {
        self.end_time.is_some()
    }

    /// 获取持续时间
    pub fn get_duration(&self) -> Option<Duration> {
        self.duration.map(|d| Duration::from_micros(d))
    }
}

/// Trace（完整的请求链路）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Trace {
    /// Trace ID
    pub trace_id: String,
    /// 根Span ID
    pub root_span_id: String,
    /// 服务名称
    pub service_name: String,
    /// 所有Span
    pub spans: Vec<Span>,
    /// 开始时间
    pub start_time: u64,
    /// 结束时间
    pub end_time: Option<u64>,
    /// 总持续时间（微秒）
    pub duration: Option<u64>,
    /// 状态
    pub status: SpanStatus,
}

impl Trace {
    pub fn new(root_span: Span) -> Self {
        let trace_id = root_span.trace_id.clone();
        let root_span_id = root_span.span_id.clone();
        let service_name = root_span.service_name.clone();
        let start_time = root_span.start_time;
        
        Self {
            trace_id,
            root_span_id,
            service_name,
            spans: vec![root_span],
            start_time,
            end_time: None,
            duration: None,
            status: SpanStatus::Unset,
        }
    }

    /// 添加Span
    pub fn add_span(&mut self, span: Span) {
        self.spans.push(span);
    }

    /// 完成Trace
    pub fn finish(mut self) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_micros() as u64;
        
        self.end_time = Some(now);
        self.duration = Some(now - self.start_time);
        
        // 计算整体状态
        self.status = if self.spans.iter().any(|s| s.attributes.status == SpanStatus::Error) {
            SpanStatus::Error
        } else {
            SpanStatus::Ok
        };
        
        self
    }

    /// 获取根Span
    pub fn get_root_span(&self) -> Option<&Span> {
        self.spans.iter().find(|s| s.span_id == self.root_span_id)
    }

    /// 获取子Span
    pub fn get_child_spans(&self, parent_id: &str) -> Vec<&Span> {
        self.spans
            .iter()
            .filter(|s| s.parent_id.as_ref() == Some(&parent_id.to_string()))
            .collect()
    }
}

/// 追踪导出器trait
#[async_trait]
pub trait TraceExporter: Send + Sync {
    /// 导出Span批次
    async fn export_spans(&self, spans: Vec<Span>) -> ServiceResult<()>;
    
    /// 导出Trace
    async fn export_trace(&self, trace: Trace) -> ServiceResult<()>;
    
    /// 关闭导出器
    async fn shutdown(&self) -> ServiceResult<()>;
}

/// 内存追踪导出器（用于测试和开发）
pub struct InMemoryTraceExporter {
    spans: Arc<RwLock<Vec<Span>>>,
    traces: Arc<RwLock<Vec<Trace>>>,
}

impl InMemoryTraceExporter {
    pub fn new() -> Self {
        Self {
            spans: Arc::new(RwLock::new(Vec::new())),
            traces: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// 获取所有导出的Span
    pub async fn get_spans(&self) -> Vec<Span> {
        self.spans.read().await.clone()
    }

    /// 获取所有导出的Trace
    pub async fn get_traces(&self) -> Vec<Trace> {
        self.traces.read().await.clone()
    }

    /// 清空所有数据
    pub async fn clear(&self) {
        self.spans.write().await.clear();
        self.traces.write().await.clear();
    }
}

#[async_trait]
impl TraceExporter for InMemoryTraceExporter {
    async fn export_spans(&self, spans: Vec<Span>) -> ServiceResult<()> {
        let mut stored_spans = self.spans.write().await;
        stored_spans.extend(spans);
        Ok(())
    }

    async fn export_trace(&self, trace: Trace) -> ServiceResult<()> {
        let mut stored_traces = self.traces.write().await;
        stored_traces.push(trace);
        Ok(())
    }

    async fn shutdown(&self) -> ServiceResult<()> {
        // 内存导出器不需要特殊的关闭逻辑
        Ok(())
    }
}

/// 控制台追踪导出器（输出到标准输出）
pub struct ConsoleTraceExporter {
    pretty_print: bool,
}

impl ConsoleTraceExporter {
    pub fn new(pretty_print: bool) -> Self {
        Self { pretty_print }
    }
}

#[async_trait]
impl TraceExporter for ConsoleTraceExporter {
    async fn export_spans(&self, spans: Vec<Span>) -> ServiceResult<()> {
        for span in spans {
            if self.pretty_print {
                println!("{}", serde_json::to_string_pretty(&span)
                    .map_err(|e| ServiceError::MetricsError(format!("序列化Span失败: {}", e)))?);
            } else {
                println!("{}", serde_json::to_string(&span)
                    .map_err(|e| ServiceError::MetricsError(format!("序列化Span失败: {}", e)))?);
            }
        }
        Ok(())
    }

    async fn export_trace(&self, trace: Trace) -> ServiceResult<()> {
        if self.pretty_print {
            println!("{}", serde_json::to_string_pretty(&trace)
                .map_err(|e| ServiceError::MetricsError(format!("序列化Trace失败: {}", e)))?);
        } else {
            println!("{}", serde_json::to_string(&trace)
                .map_err(|e| ServiceError::MetricsError(format!("序列化Trace失败: {}", e)))?);
        }
        Ok(())
    }

    async fn shutdown(&self) -> ServiceResult<()> {
        // 控制台导出器不需要特殊的关闭逻辑
        Ok(())
    }
}

/// 追踪处理器配置
#[derive(Debug, Clone)]
pub struct TracingConfig {
    /// 批处理大小
    pub batch_size: usize,
    /// 批处理超时时间（毫秒）
    pub batch_timeout_ms: u64,
    /// 最大队列大小
    pub max_queue_size: usize,
    /// 是否启用采样
    pub sampling_enabled: bool,
    /// 采样率（0.0-1.0）
    pub sampling_rate: f64,
}

impl Default for TracingConfig {
    fn default() -> Self {
        Self {
            batch_size: 100,
            batch_timeout_ms: 5000,
            max_queue_size: 10000,
            sampling_enabled: true,
            sampling_rate: 1.0,
        }
    }
}

/// 追踪处理器
pub struct TracingProcessor {
    config: TracingConfig,
    exporter: Arc<dyn TraceExporter>,
    span_queue: Arc<RwLock<Vec<Span>>>,
    trace_queue: Arc<RwLock<Vec<Trace>>>,
}

impl TracingProcessor {
    pub fn new(config: TracingConfig, exporter: Arc<dyn TraceExporter>) -> Self {
        Self {
            config,
            exporter,
            span_queue: Arc::new(RwLock::new(Vec::new())),
            trace_queue: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// 处理单个Span
    pub async fn process_span(&self, span: Span) -> ServiceResult<()> {
        // 检查采样
        if self.config.sampling_enabled && !self.should_sample() {
            return Ok(());
        }

        let mut queue = self.span_queue.write().await;
        queue.push(span);

        // 检查是否需要批量导出
        if queue.len() >= self.config.batch_size {
            let spans_to_export = queue.drain(..).collect();
            drop(queue);
            self.exporter.export_spans(spans_to_export).await?;
        }

        Ok(())
    }

    /// 处理Trace
    pub async fn process_trace(&self, trace: Trace) -> ServiceResult<()> {
        // 检查采样
        if self.config.sampling_enabled && !self.should_sample() {
            return Ok(());
        }

        let mut queue = self.trace_queue.write().await;
        queue.push(trace);

        // 检查是否需要批量导出
        if queue.len() >= self.config.batch_size {
            let traces_to_export: Vec<Trace> = queue.drain(..).collect();
            drop(queue);
            
            // 将Trace转换为Span进行导出
            for trace in traces_to_export {
                self.exporter.export_trace(trace).await?;
            }
        }

        Ok(())
    }

    /// 强制刷新所有队列
    pub async fn flush(&self) -> ServiceResult<()> {
        // 导出所有待处理的Span
        let spans = {
            let mut queue = self.span_queue.write().await;
            queue.drain(..).collect::<Vec<_>>()
        };
        if !spans.is_empty() {
            self.exporter.export_spans(spans).await?;
        }

        // 导出所有待处理的Trace
        let traces = {
            let mut queue = self.trace_queue.write().await;
            queue.drain(..).collect::<Vec<_>>()
        };
        for trace in traces {
            self.exporter.export_trace(trace).await?;
        }

        Ok(())
    }

    /// 关闭处理器
    pub async fn shutdown(&self) -> ServiceResult<()> {
        self.flush().await?;
        self.exporter.shutdown().await?;
        Ok(())
    }

    /// 判断是否应该采样
    fn should_sample(&self) -> bool {
        if !self.config.sampling_enabled {
            return true;
        }
        
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        use std::time::{SystemTime, UNIX_EPOCH};
        
        let mut hasher = DefaultHasher::new();
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos()
            .hash(&mut hasher);
        
        let hash = hasher.finish();
        let sample_threshold = (self.config.sampling_rate * u64::MAX as f64) as u64;
        hash <= sample_threshold
    }

    /// 获取队列统计信息
    pub async fn get_queue_stats(&self) -> (usize, usize) {
        let span_queue_size = self.span_queue.read().await.len();
        let trace_queue_size = self.trace_queue.read().await.len();
        (span_queue_size, trace_queue_size)
    }
}